"""PostgreSQL Connection Handler."""
import prometheus_client as prometheus
import psycopg2

METRIC_TIME = prometheus.Histogram(
    'cki_psql_query_duration_seconds', 'Time spent handling a query',
    ['hostname', 'query'],
)


class PSQLHandler:
    """PostgreSQL connection handler."""

    def __init__(self,
                 host, port, dbname, *,
                 user=None, password=None, require_ssl=False, autocommit=True):
        # pylint: disable=too-many-arguments,too-many-positional-arguments
        """
        Initialize PostgreSQL connection.

        Parameters
            - host: PostgreSQL host (required).
            - port: PostgreSQL port (required).
            - dbname: Database name (required).
            - user: Username for authentication. Default: `None`.
            - password: Password for authentication. Default: `None`.
            - require_ssl: Set sslmode=require. Default: `False`.
            - autocommit: Enable isolation level: Autocommit. Default: `True`.
        """
        self.host = host
        self.connection_string = f"host={host} port={port} dbname={dbname}"
        if user:
            self.connection_string += f" user={user}"
        if password:
            self.connection_string += f" password={password}"
        if require_ssl:
            self.connection_string += " sslmode=require"

        self.autocommit = autocommit
        self._connection = None

    @property
    def connection(self):
        """
        Get DB connection.

        Split from init for simpler testing.
        """
        if not self._connection or self._connection.closed:
            self._connection = psycopg2.connect(self.connection_string)

            if self.autocommit:
                self._connection.set_isolation_level(
                    psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT
                )

        return self._connection

    def _execute(self, query, *args):
        with METRIC_TIME.labels(self.host, query).time(), self.connection.cursor() as cursor:
            cursor.execute(query, *args)
            return cursor.fetchall()

    def execute(self, query, *args):
        """Execute a given query."""
        try:
            return self._execute(query, *args)
        except psycopg2.InterfaceError:
            # We only know the connection is closed after executing a query
            # and failing. Retry it once to avoid failing because of closed connections.
            return self._execute(query, *args)
